import os
import glob
import numpy as np
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import (
    classification_report,
    accuracy_score, precision_score, recall_score, f1_score, confusion_matrix
)

from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from xgboost import XGBClassifier
from lightgbm import LGBMClassifier
from catboost import CatBoostClassifier

# —— 用户可选模型 ——
available_models = {
    'RandomForest': RandomForestClassifier(n_estimators=100, random_state=42),
    'LogisticRegression': LogisticRegression(max_iter=1000),
    'XGBoost': XGBClassifier(use_label_encoder=False, eval_metric='logloss'),
    'LightGBM': LGBMClassifier(),
    'CatBoost': CatBoostClassifier(verbose=0)
}
# 在此列表中添加你想运行的模型名称
models_to_run = ['RandomForest', 'LogisticRegression', 'LightGBM']

# —— 定义多组特征组合 ——
feature_sets = {
    'core4':        ['rmssd', 'hf', 'pnni_50', 'lf_hf_ratio', 'sampen'],
    'time_domain':  ['sdnn','rmssd','pnni_50','pnni_20','sdsd','sampen','lf_hf_ratio'],
    'freq_domain':  ['lf_hf_ratio','vlf'],
    'nonlinear':    ['ApEn','sampen','triangular_index','lf_hf_ratio','vlf'],
    'hrv_all':      ['mean_nni','sdnn','rmssd','pnni_50','pnni_20',
                     'lf','hf','lf_hf_ratio','vlf','total_power','ApEn','sampen']
}

# 构建第3-8个window的特征平均，要求至少10个window

def build_feature_dataframe(feature_folder, min_windows=10, start_idx=2, end_idx=8):
    compliant = 0
    non_compliant = 0
    records = []
    for fp in glob.glob(os.path.join(feature_folder, '*_rrfeature.csv')):
        pid = os.path.basename(fp).replace('_rrfeature.csv', '')
        df = pd.read_csv(fp)
        win = df[df['label']=='window']
        if len(win) < min_windows:
            non_compliant += 1
            continue
        subset = win.iloc[start_idx:end_idx]
        rec = {'patient_id': pid}
        for col in set(sum(feature_sets.values(), [])):
            vals = pd.to_numeric(subset[col], errors='coerce') if col in subset else pd.Series([])
            rec[col] = vals.mean() if not vals.dropna().empty else np.nan
        compliant += 1
        records.append(rec)
    print(f"✅ 合规样本: {compliant}, 不足 {min_windows} 窗口: {non_compliant}")
    return pd.DataFrame(records)

# 读取患者性别与年龄名录，并可筛选年龄

def load_gender_info(roster_csv_path, max_age=None):
    df = pd.read_csv(roster_csv_path, dtype=str)
    # 合成 patient_id
    df['patient_id'] = df['userid'].str.strip() + '_' + df['recordid'].str.strip()
    df['gender'] = df['Gender'].str.strip()
    df['age'] = pd.to_numeric(df['age'], errors='coerce')
    if max_age is not None:
        df = df[df['age'] < max_age]
        print(f"✅ 筛选年龄<{max_age} 的样本: {len(df)} 位")
    return df[['patient_id','gender']]

# 多模型训练、打印报告并保存结果

def train_and_save(feature_folder, roster_csv, min_windows=10, max_age=None, out_csv=None):
    feat_df = build_feature_dataframe(feature_folder, min_windows)
    info_df = load_gender_info(roster_csv, max_age)
    merged = feat_df.merge(info_df, on='patient_id', how='inner').dropna(subset=['gender'])

    le = LabelEncoder()
    le.fit(merged['gender'])
    y_map = dict(zip(merged['patient_id'], le.transform(merged['gender'])))

    # 可选模型
    classifiers = {name: model for name, model in available_models.items() if name in models_to_run}

    results = []
    for set_name, feats in feature_sets.items():
        X = merged[['patient_id'] + feats].set_index('patient_id')
        X = X.replace([np.inf, -np.inf], np.nan).dropna()
        if X.empty:
            print(f"⚠️ 特征集 {set_name} 无有效样本，跳过")
            continue
        y = X.index.map(y_map)

        X_tr, X_te, y_tr, y_te = train_test_split(
            X.astype(float), list(y), test_size=0.3,
            random_state=42, stratify=list(y)
        )

        for name, model in classifiers.items():
            model.fit(X_tr, y_tr)
            y_pred = model.predict(X_te)

            print(f"\n=== 特征集：{set_name}，模型：{name} ===")
            print(classification_report(y_te, y_pred, target_names=le.classes_, zero_division=0))

            cm = confusion_matrix(y_te, y_pred)
            tn, fp, fn, tp = cm.ravel() if cm.shape==(2,2) else (np.nan,)*4
            metrics = {
                'feature_set': set_name,
                'model': name,
                'accuracy': accuracy_score(y_te, y_pred),
                'precision': precision_score(y_te, y_pred, zero_division=0),
                'recall': recall_score(y_te, y_pred, zero_division=0),
                'f1': f1_score(y_te, y_pred, zero_division=0),
                'TN': tn, 'FP': fp, 'FN': fn, 'TP': tp
            }
            results.append({**metrics, 'type':'metrics'})

            if hasattr(model, 'feature_importances_'):
                imps = model.feature_importances_
            elif hasattr(model, 'coef_'):
                imps = np.abs(model.coef_[0])
            else:
                imps = [np.nan]*len(feats)
            for feat, imp in zip(feats, imps):
                results.append({
                    'feature_set': set_name,
                    'model': name,
                    'type': 'importance',
                    'feature': feat,
                    'importance': imp
                })

    df_res = pd.DataFrame(results)
    out_csv = out_csv or 'feature_combination_results.csv'
    df_res.to_csv(out_csv, index=False)
    print(f"\n✅ 全部结果已保存到：{out_csv}")

if __name__ == '__main__':
    feature_folder = '/database/home/duansizhang/hrv_predict/data/rr_featrue/'
    roster_csv     = '/database/private/mgcdb/info_xml_raw_match_1.csv'
    # max_age 可设为 50 以筛选 50 岁以下患者，不需筛选设为 None
    train_and_save(
        feature_folder,
        roster_csv,
        min_windows=10,
        max_age=50,
        out_csv='/database/home/duansizhang/hrv_predict/result/combination_results.csv'
    )

